# -*- coding: utf-8 -*-
import asyncio
import aiohttp
import json
from globalvariables import *
from zhipuai import ZhipuAI
import time


class Agent:
    def __init__(self, token):
        self.token = token

    def chat(self, user_input: str,
             response_format_type: str = "text",
             language: str = "",
             json_format: str = "",
             temperature: float = 0.95,
             top_p: float = 0.7,
             model_code: str = "glm-4-plus",
             max_tokens: int = 1024,
             do_sample: bool = True):
        """

        :param user_input:
        :param response_format_type:
        :param language:
        :param json_format:
        :param temperature: 控制生成的随机性，值越高（如1.0），输出越随机；值越低（如0.2），输出越确定性
                   0.7 是一个常用值，通常表现较好，适用于大部分场景。
        :param top_p: 核采样参数，控制生成的多样性，值越低生成越集中于概率最高的选项。值通常设为 0.9，意味着取前 90% 概率的词作为候选
        :return:
        """
        # 选择系统提示词语言
        bg_system_prompt = bg_system_prompts["English"]
        json_system_prompt_basic = json_system_prompts_basic["English"]
        json_format_prompt = json_format_prompts["English"]
        if language == "Chinese":
            bg_system_prompt = bg_system_prompts["Chinese"]
            json_system_prompt_basic = json_system_prompts_basic["Chinese"]
            json_format_prompt = json_format_prompts["Chinese"]

        # 系统提示词
        system_prompts = [bg_system_prompt] # 首先将背景系统提示词加入到system_prompt中
        if response_format_type == "json_object":
            json_system_prompt = json_system_prompt_basic
            if json_format != "":
                json_system_prompt = json_system_prompt + json_format_prompt + json_format
            system_prompts.append(json_system_prompt)

        if language != "":
            if language == "Chinese":
                language_system_prompt = "将输出结果语种设置为中文"
            else:
                language_system_prompt = f"Set the output language to {language}"
            system_prompts.append(language_system_prompt)

        # 将该次任务需要的Prompt加入到system_prompt当中
        messages = [{"role": "system", "content": system_prompt} for system_prompt in system_prompts]
        messages.append({"role": "user", "content": user_input})

        # 初始化ZhipuAI
        client = ZhipuAI(api_key=self.token)

        response = client.chat.completions.create(
            model=model_code,  # 要调用的模型名称
            messages=messages,
            response_format={"type": response_format_type},
            temperature=temperature,
            top_p=top_p,
            max_tokens=max_tokens,
            do_sample=do_sample
        )
        return response.choices[0].message.content

    async def chat_async(self, user_input: str,
                         response_format_type: str = "text",
                         language: str = "",
                         json_format: str = "",
                         temperature: float = 0.95,
                         top_p: float = 0.7,
                         model_code: str = "glm-4-plus",
                         max_tokens: int = 1024,
                         do_sample: bool = True):

        """
                异步调用ZhipuAI进行聊天对话生成。

                :param user_input: 用户输入的文本
                :param response_format_type: 响应格式
                :param language: 输出语言
                :param json_format: JSON格式
                :param temperature: 随机性控制，值越高输出越随机
                :param top_p: 核采样控制，越低生成的内容越集中
                :param model_code: 模型名称
                :param max_tokens: 最大输出长度
                :param do_sample: 是否启用采样
                :return: 模型生成的文本
                """
        # 选择系统提示词语言
        bg_system_prompt = bg_system_prompts["English"]
        json_system_prompt_basic = json_system_prompts_basic["English"]
        json_format_prompt = json_format_prompts["English"]
        if language == "Chinese":
            bg_system_prompt = bg_system_prompts["Chinese"]
            json_system_prompt_basic = json_system_prompts_basic["Chinese"]
            json_format_prompt = json_format_prompts["Chinese"]

        # 系统提示词
        system_prompts = [bg_system_prompt] # 首先将背景系统提示词加入到system_prompt中
        if response_format_type == "json_object":
            json_system_prompt = json_system_prompt_basic
            if json_format != "":
                json_system_prompt = json_system_prompt + json_format_prompt + json_format
            system_prompts.append(json_system_prompt)

        if language != "":
            if language == "Chinese":
                language_system_prompt = "将输出结果语种设置为中文"
            else:
                language_system_prompt = f"Set the output language to {language}"
            system_prompts.append(language_system_prompt)

        # 设置消息
        messages = [{"role": "system", "content": system_prompt} for system_prompt in system_prompts]
        messages.append({"role": "user", "content": user_input})
        # 异步调用API
        # client = ZhipuAI(api_key="efedb300c8468ae1315a5474228100f6.cZbs8qwDylqU2GzD")
        url = "https://open.bigmodel.cn/api/paas/v4/chat/completions"
        headers = {
            "Authorization": "Bearer " + self.token,  # API Key
            "Content-Type": "application/json"
        }
        retries = 0
        max_retries = 3
        while retries < max_retries:
            try:
                # 使用 aiohttp 进行异步请求
                async with aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=120)) as session:
                    # 构建请求体
                    body = {
                        "model": model_code,  # 要调用的模型名称
                        "messages": messages,
                        "temperature": temperature,
                        "top_p": top_p,
                        "max_tokens": max_tokens,
                        "do_sample": do_sample
                    }
                    # 发送 POST 请求
                    async with session.post(url, json=body, headers=headers) as response:
                        # 确保请求成功
                        if response.status == 200:
                            response_text = await response.text()
                            if response_text:
                                result = await response.json()  # 解析 JSON 响应
                                # 获取模型生成的内容
                                return result['choices'][0]['message']['content']
                            else:
                                print("Error: Response content is empty.")
                        else:
                            print(f"Error: {response.status}, {await response.text()}")
            except aiohttp.ClientError as e:
                print(f"Network error during API call: {e}. Retrying in 3 seconds...")
                retries += 1
                if retries < max_retries:
                    await asyncio.sleep(3)
            except ValueError as e:
                print(f"Invalid JSON format: {e}")
                break
            except Exception as e:
                print(f"Unexpected error during API call: {e}")
                break
        return None
